"""
langchain 内置了RAG的功能 

"""

from langchain.document_loaders import UnstructuredMarkdownLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.llms import OpenAI
from langchain.chains import RetrievalQA
from langchain.document_loaders import PyPDFLoader
from dotenv import load_dotenv,find_dotenv
_ =load_dotenv(find_dotenv());

# 加载文档
loader = PyPDFLoader('E:\\project\\ai\\zhihu_llm\\RAG\\rzf.pdf')
pages = loader.load_and_split()

# 文档切分       递归字符文本拆分器
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50,
                                               length_function=len,add_start_index=True)
# 将pages[0].page_content,pages[1].page_content,这两页的内容切分为多个文档
text = text_splitter.create_documents([pages[0].page_content,
                                       pages[1].page_content,])

# 存入数据库
embeddings = OpenAIEmbeddings()
db  = Chroma.from_documents(text, embeddings)

# langchain内置RAG实现  
qa_chain = RetrievalQA.from_chain_type(
  llm=OpenAI(temperature=0),
  retriever =db.as_retriever()
  )

query = "任正非是干什么的?"
response = qa_chain.run(query)
print(response)